{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "W8_Tutorial2.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/main/tutorials/W08_AutoEncoders_GANs/W8_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2kPrtVyXtaP"
      },
      "source": [
        "# CIS-522 Week 8 Part 2\n",
        "# Generative Adversarial Networks (GANs)\n",
        "\n",
        "__Instructor__: Konrad Kording\n",
        "\n",
        "__Content creators:__ Arash Ash, Richard Lange"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p2DuJddRXxYY"
      },
      "source": [
        "---\n",
        "# Tutorial Objectives\n",
        "\n",
        "This tutorial will uncover Generative Adversarial Networks (GANs) and teach you how to use them to solve all kinds of cool computer vision tasks. In particular, we will:\n",
        "\n",
        "* First, understand how GANs are trained\n",
        "* Generate Animal Faces with MLPs and ConvNets\n",
        "* Discover the most common failure mode of GANs\n",
        "* How to train conditionally on the class label\n",
        "* How to generate images controllably using a pre-trained classifier \n",
        "* How to learn an image-to-image mapping using unpaired images using CycleGANs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "e6GnDC12UtUD"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time();\n",
        "\n",
        "my_pennkey = 'value' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','astute-jellyfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NtskSZXZqu3W",
        "cellView": "form"
      },
      "source": [
        "# @title Slides\n",
        "from IPython.display import HTML\n",
        "HTML('<iframe src=\"https://docs.google.com/presentation/d/e/2PACX-1vQTsDQET4TdoyJv_Ql_wm6ct9cxFMVghMXnqGYutHJW_3io7D1wM_F0cOj1-31KwQ-gV5eHWSni77mO/embed?start=false&loop=false&delayms=3000\" frameborder=\"0\" width=\"480\" height=\"299\" allowfullscreen=\"true\" mozallowfullscreen=\"true\" webkitallowfullscreen=\"true\"></iframe>')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vgOQxVS1X2dB"
      },
      "source": [
        "---\n",
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xXd8Now_XwR7"
      },
      "source": [
        "# imports\n",
        "import torch\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import time\n",
        "import random\n",
        "import pathlib\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as transforms\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from torchvision.utils import make_grid\n",
        "from IPython.display import HTML, display\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ZUyEQpxTSzxe"
      },
      "source": [
        "# @title Seeding for reproducibility\n",
        "seed = 522\n",
        "random.seed(seed)\n",
        "np.random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed(seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "hj8C2Cf0G0i_"
      },
      "source": [
        "# @title Downloads\n",
        "%%capture\n",
        "!rm -r AnimalFaces32x32/\n",
        "!git clone https://github.com/arashash/AnimalFaces32x32\n",
        "!rm -r afhq/\n",
        "!unzip ./AnimalFaces32x32/afhq_32x32.zip \n",
        "\n",
        "!git clone https://github.com/arashash/celeb-classifier\n",
        "!mv /content/celeb-classifier/pretrained_classifier.pth pretrained_classifier.pth"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "rr1VWDwRyraD"
      },
      "source": [
        "# @title Animal Faces Dataset\n",
        "\n",
        "def get_data_loaders(batch_size):\n",
        "  # define the transform done only during training\n",
        "  augmentation_transforms = [transforms.RandomHorizontalFlip()]\n",
        "\n",
        "  # define the transform done in training and testing (after augmentation)\n",
        "  preprocessing_transforms = [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
        "  \n",
        "  # compose them together\n",
        "  train_transform = transforms.Compose(augmentation_transforms + preprocessing_transforms)\n",
        "  test_transform = transforms.Compose(preprocessing_transforms)\n",
        "\n",
        "  # using pathlib to be compatible with all OS's\n",
        "  data_path = pathlib.Path('.')/'afhq'\n",
        "\n",
        "  # define the dataset objects (they can load one by one)\n",
        "  img_train_dataset = ImageFolder(data_path/'train', transform=train_transform)\n",
        "  img_test_dataset = ImageFolder(data_path/'val', transform=test_transform)\n",
        "\n",
        "  # define the dataloader objects (they can load batch by batch)\n",
        "  img_train_loader = DataLoader(img_train_dataset, batch_size=batch_size, pin_memory=True,\n",
        "                                shuffle=True, num_workers=0, drop_last=True)\n",
        "  # num_workers can be set to higher if running on Colab Pro TPUs to speed up,\n",
        "  # with more than one worker, it will do multithreading to queue batches\n",
        "  img_test_loader = DataLoader(img_test_dataset, batch_size=batch_size,\n",
        "                          shuffle=False, num_workers=0)\n",
        "  \n",
        "  return img_train_loader, img_test_loader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "oRrTcZOlXA8C"
      },
      "source": [
        "# @title Figure Settings\n",
        "import ipywidgets as widgets\n",
        "%matplotlib inline \n",
        "fig_w, fig_h = (8, 6)\n",
        "plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "my_layout = widgets.Layout()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "40Qc7MFcX7Oc"
      },
      "source": [
        "# @title Helper functions\n",
        "def imshow(img):\n",
        "    img -= torch.min(img)\n",
        "    img *= 255/torch.max(img)\n",
        "    npimg = img.numpy().astype(np.uint8)\n",
        "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
        "    plt.axis(False)\n",
        "    plt.show()\n",
        "\n",
        "def show_tensor_images(image_tensor, num_images=16, size=(3, 64, 64), nrow=3):\n",
        "    image_tensor = (image_tensor + 1) / 2\n",
        "    image_unflat = image_tensor.detach().cpu()\n",
        "    image_grid = make_grid(image_unflat[:num_images], nrow=nrow)\n",
        "    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n",
        "    plt.show()\n",
        "\n",
        "def progress(batch, loss, batches):\n",
        "    return HTML(\"\"\"\n",
        "        <label for=\"file\">Training loss: {loss}</label>\n",
        "        <progress\n",
        "            value='{batch}'\n",
        "            max='{batches}',\n",
        "            style='width: 100%'\n",
        "        >\n",
        "            {batch}\n",
        "        </progress>\n",
        "    \"\"\".format(loss=loss, batch=batch, batches=batches))\n",
        "\n",
        "class Classifier(nn.Module):\n",
        "    def __init__(self, im_chan=3, n_classes=2, hidden_dim=64):\n",
        "        super(Classifier, self).__init__()\n",
        "        self.classifier = nn.Sequential(\n",
        "            self.make_classifier_block(im_chan, hidden_dim),\n",
        "            self.make_classifier_block(hidden_dim, hidden_dim * 2),\n",
        "            self.make_classifier_block(hidden_dim * 2, hidden_dim * 4, stride=3),\n",
        "            self.make_classifier_block(hidden_dim * 4, n_classes, final_layer=True),\n",
        "        )\n",
        "\n",
        "    def make_classifier_block(self, input_channels, output_channels, kernel_size=4, stride=2, final_layer=False):\n",
        "        if final_layer:\n",
        "            return nn.Sequential(\n",
        "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
        "            )\n",
        "        else:\n",
        "            return nn.Sequential(\n",
        "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
        "                nn.BatchNorm2d(output_channels),\n",
        "                nn.LeakyReLU(0.2, inplace=True),\n",
        "            )\n",
        "\n",
        "    def forward(self, image):\n",
        "        class_pred = self.classifier(image)\n",
        "        return class_pred.view(len(class_pred), -1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SG3iPDUaxSLv"
      },
      "source": [
        "---\n",
        "# Section 1: wrapping up VAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "CwszyupvkL3-"
      },
      "source": [
        "#@title Video: Tutorial 1 Recap\n",
        "try: t1;\n",
        "except NameError: t1=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"F5FBVoQKtAY\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MtToppuEAjGk"
      },
      "source": [
        "Before anything, let's memorize these global variables for the rest of the tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m7Ku8MrOAiFe"
      },
      "source": [
        "############################ Constants #############################\n",
        "H = 32 # image height\n",
        "W = 32 # image width\n",
        "C = 3 # image channel size, e.g., 3 for RGB\n",
        "num_classes = 3 # image labels; cat, dot, wild\n",
        "batch_size = 32 # dataloaders batch size\n",
        "\n",
        "# noise dimension for the genereator, i.e., latent variable\n",
        "# drawn from standard normal distribution\n",
        "z_dim = 128\n",
        "####################################################################"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yRsWs55XfyAS"
      },
      "source": [
        "## Exercise 1: why vanilla VAEs fail?\n",
        "\n",
        "Now that we have a pretty good idea about Auto-encoders let's see why they fail in practice? Let's train a VAE on the Animal Faces dataset to see whether we can generate cute cat faces?\n",
        "\n",
        "first, remembering the Animal Faces dataset:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2IT6DryFXpaV"
      },
      "source": [
        "# creating the dataloaders which gives us batches of samples\n",
        "img_train_loader, img_test_loader = get_data_loaders(batch_size)\n",
        "\n",
        "# get some random training images\n",
        "dataiter = iter(img_train_loader)\n",
        "images, labels = dataiter.next()\n",
        "\n",
        "# show images\n",
        "imshow(make_grid(images, nrow=8))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "MFcrKW-ZBQQh"
      },
      "source": [
        "# @markdown **Run me** (a copy of `rsample`, `log_p_x`, `kl_q_p`, `BiasLayer`, and `ConvVAE` from Tutorial1)\n",
        "my_dataset_size = C, H, W\n",
        "def rsample(phi, n_samples):\n",
        "    \"\"\"Sample z ~ q(z;phi)\n",
        "    Ouput z is size [b,n_samples,K] given phi with shape [b,K+1]. The first K\n",
        "    entries of each row of phi are the mean of q, and phi[:,-1] is the log\n",
        "    standard deviation\n",
        "    \"\"\"\n",
        "    b, kplus1 = phi.size()\n",
        "    eps = torch.randn(b, n_samples, kplus1-1, device=phi.device)\n",
        "    mu, sig = phi[:, :-1], phi[:,-1].exp()\n",
        "    return eps*sig.view(b,1,1) + mu.view(b,1,kplus1-1)\n",
        "\n",
        "def log_p_x(x, mu_xs, sig_x):\n",
        "    \"\"\"Given [batch, ...] input x and [batch, n, ...] reconstructions, compute\n",
        "    pixel-wise log Gaussian probability\n",
        "\n",
        "    Sum over pixel dimensions, but mean over batch and samples.\n",
        "    \"\"\"\n",
        "    b, n = mu_xs.size()[:2]\n",
        "    # Flatten out pixels and add a singleton dimension [1] so that x will be\n",
        "    # implicitly expanded when combined with mu_xs\n",
        "    x = x.reshape(b, 1, -1)\n",
        "    _, _, p = x.size()\n",
        "    squared_error = (x - mu_xs.view(b, n, -1))**2 / (2*sig_x**2)\n",
        "\n",
        "    # Size of squared_error is [b,n,p]. log prob is by definition sum over [p].\n",
        "    # Expected value requires mean over [n]. Handling different size batches\n",
        "    # requires mean over [b].\n",
        "    return -(squared_error + torch.log(sig_x)).sum(dim=2).mean(dim=(0,1))\n",
        "\n",
        "def kl_q_p(zs, phi):\n",
        "    \"\"\"Given [b,n,k] samples of z drawn from q, compute estimate of KL(q||p).\n",
        "    phi must be size [b,k+1]\n",
        "    \"\"\"\n",
        "    b, n, k = zs.size()\n",
        "    mu_p = torch.zeros((k,), device=phi.device)\n",
        "    sig_p = torch.ones((k,), device=phi.device)\n",
        "    mu_q, log_sig_q = phi[:,:-1], phi[:,-1]\n",
        "    log_p = -0.5*(zs - mu_p.view(1,1,k))**2 / sig_p.view(1,1,k)**2 - torch.log(sig_p.view(1,1,k))\n",
        "    log_q = -0.5*(zs - mu_q.view(b,1,k))**2 / log_sig_q.exp().view(b,1,1)**2 - log_sig_q.view(b,1,-1)\n",
        "    # Size of log_q and log_p is [b,n,k]. Sum along [k] but mean along [b,n]\n",
        "    return (log_q - log_p).sum(dim=2).mean(dim=(0,1))\n",
        "\n",
        "class BiasLayer(nn.Module):\n",
        "    def __init__(self, shape):\n",
        "        super(BiasLayer, self).__init__()\n",
        "        init_bias = torch.zeros(shape)\n",
        "        self.bias = nn.Parameter(init_bias, requires_grad=True)\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return x + self.bias\n",
        "\n",
        "class ConvVAE(nn.Module):\n",
        "    def __init__(self, K, num_filters=32, filter_size=5):\n",
        "        super(ConvVAE, self).__init__()\n",
        "        \n",
        "        # With padding=0, the number of pixels cut off from each image dimension\n",
        "        # is filter_size // 2. Double it to get the amount of pixels lost in\n",
        "        # width and height per Conv2D layer, or added back in per \n",
        "        # ConvTranspose2D layer.\n",
        "        filter_reduction = 2 * (filter_size // 2)\n",
        "\n",
        "        # After passing input through two Conv2d layers, the shape will be\n",
        "        # 'shape_after_conv'. This is also the shape that will go into the first\n",
        "        # deconvolution layer in the decoder\n",
        "        self.shape_after_conv = (num_filters,\n",
        "                                 my_dataset_size[1]-2*filter_reduction,\n",
        "                                 my_dataset_size[2]-2*filter_reduction)\n",
        "        flat_size_after_conv = self.shape_after_conv[0] \\\n",
        "            * self.shape_after_conv[1] \\\n",
        "            * self.shape_after_conv[2]\n",
        "\n",
        "        # Define the recognition model (encoder or q) part\n",
        "        self.q_bias = BiasLayer(my_dataset_size)      \n",
        "        self.q_conv_1 = nn.Conv2d(my_dataset_size[0], num_filters, 5)\n",
        "        self.q_conv_2 = nn.Conv2d(num_filters, num_filters, 5)\n",
        "        self.q_fc_phi = nn.Linear(flat_size_after_conv, K+1)\n",
        "\n",
        "        # Define the generative model (decoder or p) part\n",
        "        self.p_fc_upsample = nn.Linear(K, flat_size_after_conv)\n",
        "        self.p_deconv_1 = nn.ConvTranspose2d(num_filters, num_filters, 5)\n",
        "        self.p_deconv_2 = nn.ConvTranspose2d(num_filters, my_dataset_size[0], 5)\n",
        "        self.p_bias = BiasLayer(my_dataset_size)\n",
        "\n",
        "        # Define a special extra parameter to learn scalar sig_x for all pixels\n",
        "        self.log_sig_x = nn.Parameter(torch.zeros(()))\n",
        "    \n",
        "    def infer(self, x):\n",
        "        \"\"\"Map (batch of) x to (batch of) phi which can then be passed to\n",
        "        rsample to get z\n",
        "        \"\"\"\n",
        "        s = self.q_bias(x)\n",
        "        s = F.relu(self.q_conv_1(s))\n",
        "        s = F.relu(self.q_conv_2(s))\n",
        "        flat_s = s.view(s.size()[0], -1)\n",
        "        phi = self.q_fc_phi(flat_s)\n",
        "        return phi\n",
        "\n",
        "    def generate(self, zs):\n",
        "        \"\"\"Map [b,n,k] sized samples of z to [b,n,p] sized images\n",
        "        \"\"\"\n",
        "        # Note that for the purposes of passing through the generator, we need\n",
        "        # to reshape zs to be size [b*n,k]\n",
        "        b, n, k = zs.size()\n",
        "        s = zs.view(b*n, -1)\n",
        "        s = F.relu(self.p_fc_upsample(s)).view((b*n,) + self.shape_after_conv)\n",
        "        s = F.relu(self.p_deconv_1(s))\n",
        "        s = self.p_deconv_2(s)\n",
        "        s = self.p_bias(s)\n",
        "        mu_xs = s.view(b, n, -1)\n",
        "        return mu_xs\n",
        "    \n",
        "    def forward(self, x):\n",
        "        # VAE.forward() is not used for training, but we'll treat it like a\n",
        "        # classic autoencoder by taking a single sample of z ~ q\n",
        "        phi = self.infer(x)\n",
        "        zs = rsample(phi, 1)\n",
        "        return self.generate(zs).view(x.size())\n",
        "\n",
        "    def elbo(self, x, n=1):\n",
        "        # Compute ELBO for a batch of inputs\n",
        "        phi = self.infer(x)\n",
        "        zs = rsample(phi, n)\n",
        "        mu_xs = self.generate(zs)\n",
        "        return log_p_x(x, mu_xs, self.log_sig_x.exp()) - kl_q_p(zs, phi)\n",
        "\n",
        "def train_vae(vae, dataset, epochs=10, n_samples=16):\n",
        "    opt = torch.optim.Adam(vae.parameters(), lr=0.001, weight_decay=1e-6)\n",
        "    elbo_vals = []\n",
        "    vae.to(device)\n",
        "    vae.train()\n",
        "    loader = DataLoader(dataset, batch_size=100, shuffle=True, pin_memory=True)\n",
        "    for epoch in trange(epochs, desc='Epochs'):\n",
        "        for im, _ in tqdm(loader, total=len(dataset)//100, desc='Batches', leave=False):\n",
        "            im = im.to(device)\n",
        "            opt.zero_grad()\n",
        "            loss = -vae.elbo(im)\n",
        "            loss.backward()\n",
        "            opt.step()\n",
        "\n",
        "            elbo_vals.append(-loss.item())\n",
        "    vae.to('cpu')\n",
        "    vae.eval()\n",
        "    return elbo_vals"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wT1MfHWrXN51"
      },
      "source": [
        "def train_vae(vae, loader, epochs=5, n_samples=1):\n",
        "    opt = torch.optim.Adam(vae.parameters(), lr=0.001, weight_decay=1e-5)\n",
        "    \n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "\n",
        "    elbo_vals = []\n",
        "    vae.to(device)\n",
        "    vae.train()\n",
        "    steps = 0\n",
        "    for epoch in range(epochs):\n",
        "        for im, _ in loader:\n",
        "            opt.zero_grad()\n",
        "            loss = -vae.elbo(im.to(device), n_samples)\n",
        "            loss.backward()\n",
        "            opt.step()\n",
        "\n",
        "            elbo_vals.append(-loss.item())\n",
        "            progress_bar.update(progress(steps, loss.item(), max_steps))\n",
        "            steps += 1\n",
        "\n",
        "    vae.to('cpu')\n",
        "    vae.eval()\n",
        "    return elbo_vals\n",
        "\n",
        "vae = ConvVAE(K=z_dim)\n",
        "\n",
        "elbo_vals = train_vae(vae, img_train_loader, n_samples=10)\n",
        "zs = torch.randn(batch_size, 1, z_dim)\n",
        "images = vae.generate(zs).detach().view(batch_size, C, H, W)\n",
        "imshow(make_grid(images, nrow=8))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CA7r3C3ij6dx",
        "cellView": "form"
      },
      "source": [
        "#@markdown What's wrong with these generated images?\n",
        "w8_why_different_distribution = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ec8ifHJn8R2f"
      },
      "source": [
        "**Answer:** Here are the main problems,\n",
        "* Blurry\n",
        "* Not properly colored\n",
        "* All three classes are not distinguishable"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SyJYW2h6hg7w"
      },
      "source": [
        "---\n",
        "# Section 2: Intro to GANs\n",
        "*Estimated time to here: 15min*\n",
        "\n",
        "As opposed to Auto-encoders, GANs show us an intuitive way of learning representations that can be easily sampled. Let's see how,"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Nz6JNCOCjz2g"
      },
      "source": [
        "#@title Video: Intro to GANs\n",
        "try: t2;\n",
        "except NameError: t2=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"ye5P78eKOVM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Glt972-YlOoU"
      },
      "source": [
        "Let's first review the code from Week 3 (MLPs). We construct a model that can be defined both as a generator and a discriminator.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yAmVsPS7mP_E"
      },
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, actv, input_units, hidden_units_list, output_units,\n",
        "                 out_actv=None):\n",
        "        super(MLP, self).__init__()\n",
        "\n",
        "        # Initialize layers of MLP\n",
        "        self.layers = nn.ModuleList()\n",
        "\n",
        "        # Loop over layers and create each one\n",
        "        in_units = input_units\n",
        "        for units in hidden_units_list:\n",
        "          out_units = units\n",
        "\n",
        "          # create the linear layer\n",
        "          self.layers += [nn.Linear(in_units, out_units)]\n",
        "          in_units = out_units  \n",
        "\n",
        "          # create the activation layer\n",
        "          exec('self.activation = nn.%s'%actv)  \n",
        "          self.layers += [self.activation]\n",
        "\n",
        "        # Create final linear layer\n",
        "        self.out = nn.Linear(in_units, output_units)\n",
        "\n",
        "        # create output activation layer if exist\n",
        "        self.out_actv = out_actv\n",
        "        if out_actv:\n",
        "          exec('self.out_actv = nn.%s'%out_actv) \n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        # Flatten inputs to 2D (if more than that)\n",
        "        x = x.view(x.shape[0], -1)  \n",
        "\n",
        "        # Get activations of each layer\n",
        "        for layer in self.layers:\n",
        "          x = layer(x)\n",
        "\n",
        "        # Get outputs\n",
        "        x = self.out(x) \n",
        "        \n",
        "        # apply output activation if exist\n",
        "        if self.out_actv:\n",
        "          x = self.out_actv(x)\n",
        "\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eBZMcZmY6WIT"
      },
      "source": [
        "## Exercise 2.1: Discriminator Loss\n",
        "Let's implement the discriminator loss function that is introduced in the video, which is:\n",
        "\n",
        "$$\n",
        "J_{D}=-\\frac{1}{m} \\sum_{i=1}^{m} y_{i} \\log D\\left(x_{i}\\right)+\\left(1-y_{i}\\right) \\log \\left(1-D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "Note that it is often the standard to calculate this loss in two steps using equal batches ($m=2m'$) of only fake and only real examples. So let's rewrite it accordingly,\n",
        "\n",
        "$$\n",
        "J_{D}=-\\frac{1}{2m'} \\sum_{i=1}^{real} \\log D\\left(x_{i}\\right)-\\frac{1}{2m'}\\sum_{i=1}^{fake}\\log \\left(1-D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "Note that $log$ function goes to infinity near zero, which would make the loss infinite. For this reason, `nn.BCELoss` clamps its log function outputs to be greater than or equal to -100. This way, we can always have a finite loss value and a linear, backward method. we can achieve this by using `torch.clamp(x, -100, 0)` after all $log$ computations.\n",
        "\n",
        "Hint: You will need to use these functions in this order:\n",
        "\n",
        "1.   `torch.log(x)`\n",
        "2.   `torch.clamp(x, -100, 0)`\n",
        "3.   `torch.mean(x)`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tqr3Flfxf43_"
      },
      "source": [
        "def get_disc_BCE_loss(disc_real_pred, disc_fake_pred):\n",
        "    '''\n",
        "    Function for returning a BCE loss of discriminator\n",
        "\n",
        "    Parameters:\n",
        "        disc_real_pred: probs of real predicted by the disc using a real batch, (num_samples, 1)\n",
        "        disc_fake_pred: probs of real predicted by the disc using a fake batch, (num_samples, 1)\n",
        "\n",
        "    Returns:\n",
        "        BCE loss of the two prediction batches, a scalar\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the disc BCE loss function\")\n",
        "    ####################################################################\n",
        "\n",
        "    # BCE loss of the real batch\n",
        "    disc_real_loss = ...\n",
        "\n",
        "    # BCE loss of the fake batch\n",
        "    disc_fake_loss = ...\n",
        "\n",
        "    # Average loss\n",
        "    disc_loss = ...\n",
        "    return disc_loss\n",
        "\n",
        "\n",
        "def get_disc_BCE_loss(disc_real_pred, disc_fake_pred):\n",
        "    '''\n",
        "    Function for returning a BCE loss of discriminator\n",
        "\n",
        "    Parameters:\n",
        "        disc_real_pred: probs of real predicted by the disc using a real batch, (num_samples, 1)\n",
        "        disc_fake_pred: probs of real predicted by the disc using a fake batch, (num_samples, 1)\n",
        "\n",
        "    Returns:\n",
        "        BCE loss of the two prediction batches, a scalar\n",
        "    '''\n",
        "\n",
        "    # BCE loss of the real batch\n",
        "    disc_real_loss = -torch.mean(torch.clamp(torch.log(disc_real_pred), -100, 0))\n",
        "\n",
        "    # BCE loss of the fake batch\n",
        "    disc_fake_loss = -torch.mean(torch.clamp(torch.log(1 - disc_fake_pred), -100, 0))\n",
        "\n",
        "    # Average loss\n",
        "    disc_loss = (disc_real_loss + disc_fake_loss) / 2\n",
        "    return disc_loss\n",
        "\n",
        "### Uncomment below to test your function\n",
        "test_value = get_disc_BCE_loss(torch.tensor([0.9, 0.9]),\n",
        "                              torch.tensor([0.1, 0.2]))\n",
        "true_value = 0.135\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WOb-ZLWF3JZg"
      },
      "source": [
        "Before we start the training, we need the data and model. Let's define our data loaders and the MLP model for discriminators.\n",
        "\n",
        "We will first define an MLP using two hidden layers and a single Sigmoid activated output unit that predicts the probability of realness of inputs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J5WRANwB2u1S"
      },
      "source": [
        "# PReLU is like LeakyRelU\n",
        "# but the negative slope is a trainalble parameter\n",
        "actv = 'PReLU()'\n",
        "\n",
        "disc = MLP(actv=actv,\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [256, 128],\n",
        "          output_units = 1,\n",
        "          out_actv = 'Sigmoid()').to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l9YVY4fQGj9W"
      },
      "source": [
        "## Step 1: Train a Discriminator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ezXqzsudpi7d"
      },
      "source": [
        "Now we are going to train the discriminator using this loss function on random fake and real images. We will generate fake images where their pixel values are drawn from a standard normal distribution (i.e., white noise)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Suhf0gK7pzlT"
      },
      "source": [
        " def train_discriminator(disc, loader, epochs=1):\n",
        "    disc_opt = optim.Adam(disc.parameters(), lr=2e-4)\n",
        "\n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "    disc_losses = []\n",
        "    steps = 0\n",
        "    for epoch in range(epochs):\n",
        "        for i, (real, _) in enumerate(loader):\n",
        "\n",
        "          real = real.to(device)\n",
        "\n",
        "          fake = torch.randn(real.shape).to(device)\n",
        "\n",
        "          ### Train Discriminator ###\n",
        "          disc_real_pred = disc(real)\n",
        "          disc_fake_pred = disc(fake)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          disc_loss = get_disc_BCE_loss(disc_real_pred, disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          disc_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          disc_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          disc_opt.step()\n",
        "\n",
        "          # update the loss trackers\n",
        "          disc_losses.append(disc_loss.item())\n",
        "          \n",
        "          # update the progress bar\n",
        "          steps += 1\n",
        "          progress_bar.update(progress(steps, disc_losses[-1], max_steps))\n",
        "    \n",
        "    plt.plot(disc_losses, label='Discriminator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "    return disc\n",
        "\n",
        "disc_white_noise = train_discriminator(disc, img_train_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AKRaMopEiupI"
      },
      "source": [
        "fake = torch.randn(1, C*H*W).to(device)\n",
        "disc_fake_pred = disc(fake)\n",
        "'Probability of real on the fake example is %.3f'%disc_fake_pred.item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0arwPcMocTx7",
        "cellView": "form"
      },
      "source": [
        "#@markdown Add 1.0 to the fake sample generated above; does the trained discriminator classify it correctly? Why?\n",
        "w8_does_it_work_with_another_white_noise= '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iRjetIL6CL-U"
      },
      "source": [
        "**Answer:** No, it's failing miserably. So it shows that the discriminator is very sensitive to the mean of the distribution. The reason is that the real examples have very different mean than the seen fake examples, which made the discriminator task be sufficient to compare the means."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iudv_kc4hMlA"
      },
      "source": [
        "I hope you could see that distinguishing only between real and white noise is not enough. The discriminator needs to differentiate between all possible fake distributions. Therefore there should be a smart process that able to generate these fake distributions. Let's see how,"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pghLyROzGBfY"
      },
      "source": [
        "## Step 2: Train a Generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "lh9OyNUDbSCB"
      },
      "source": [
        "#@title Video: The GAN's Generator\n",
        "try: t3;\n",
        "except NameError: t3=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"wL2xy9HTcLc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOpO3g2bcJG9"
      },
      "source": [
        "## Exercise 2.2: Generator Loss\n",
        "\n",
        "Here's the BCE generator loss,\n",
        "$$\n",
        "J_{G}=-J_{D}=\\frac{1}{m} \\sum_{i=1}^{m} y_{i} \\log D\\left(x_{i}\\right)+\\left(1-y_{i}\\right) \\log \\left(1-D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "But we should get it in better shape to implement it as we did for the discriminator. First, note that the generator will only be trained with fake examples (since it can not predict real examples). Therefore the first term disappears for the generator loss since $y_i=0$. So we have,\n",
        "\n",
        "$$\n",
        "J_{G}=\\frac{1}{m} \\sum_{i=1}^{fake} \\log \\left(1-D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "<img align=\"right\" width=\"400\" height=\"250\" src=\"https://github.com/CIS-522/course-content/raw/main/tutorials/W08_AutoEncoders_GANs/static/gen_loss.png\">\n",
        "\n",
        "Furthermore, the following modified loss (proposed in the original [paper](https://arxiv.org/pdf/1406.2661.pdf)) prevents vanishing gradient in the early stages of training when the discriminator job is too easy. It achieves this goal by providing a larger gradient for the generator in the beginning compared to the BCE implementation. The tangent lines in the figure indicate the gradient at the beginning of the training where the generator is producing easy examples for the discriminator, and it has a low probability of being real. Therefore these large slopes in the new loss prevent vanishing gradient at the beginning of the training.\n",
        "\n",
        "$$\n",
        "J_{G}=-\\frac{1}{m} \\sum_{i=1}^{fake} \\log \\left(D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "This generator loss can also be implemented with BCE if we switch labels from 0 to 1, i.e., considering fake examples labeled as 1 for the generator. And in practice, this is how we use it with `nn.BCELoss`. And we will implement this way in the last exercise as well. But first, let's practice doing it ourselves from scratch!\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z-FnGEqObrSM"
      },
      "source": [
        "def get_gen_modified_BCE_loss(disc_fake_pred):\n",
        "    '''\n",
        "    Function for returning a modified BCE loss for generator\n",
        "\n",
        "    Parameters:\n",
        "        probs of real predicted by the disc\n",
        "        using a fake batch, (num_samples, 1)\n",
        "\n",
        "    Returns:\n",
        "        modified BCE loss of the fake prediction batch, a scalar\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the disc BCE loss function\")\n",
        "    ####################################################################\n",
        "\n",
        "    # modified BCE loss of the fake batch\n",
        "    gen_fake_loss = ...\n",
        "\n",
        "    return gen_fake_loss\n",
        "\n",
        "\n",
        "def get_gen_modified_BCE_loss(disc_fake_pred):\n",
        "    '''\n",
        "    Function for returning a modified BCE loss for generator\n",
        "\n",
        "    Parameters:\n",
        "        probs of real predicted by the disc\n",
        "        using a fake batch, (num_samples, 1)\n",
        "\n",
        "    Returns:\n",
        "        modified BCE loss of the fake prediction batch, a scalar\n",
        "    '''\n",
        "    \n",
        "    # modified BCE loss of the fake batch\n",
        "    gen_fake_loss = -torch.mean(torch.clamp(torch.log(disc_fake_pred), -100, 0))\n",
        "\n",
        "    return gen_fake_loss\n",
        "\n",
        "### Uncomment below to test your function\n",
        "test_value = get_gen_modified_BCE_loss(torch.tensor([0.9, 0.9]))\n",
        "true_value = 0.105\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5JD7Put5aAbf"
      },
      "source": [
        "Before we start the training, we need the generator model.\n",
        "\n",
        "Let's define an MLP using two hidden layers, and a single Tanh activated output unit that generates images given a noise vector of size `z_dim`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "liyg7h0qbSjd"
      },
      "source": [
        "gen = MLP(actv='PReLU()',\n",
        "          input_units = z_dim,\n",
        "          hidden_units_list = [128, 256],\n",
        "          output_units = C*H*W,\n",
        "          out_actv = 'Tanh()').to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xrx1_7m1ra-m"
      },
      "source": [
        "Now let's train the generator to fool the white noise discriminator we trained before:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2tIc5_lHezs4"
      },
      "source": [
        " def train_generator(disc, gen, loader, epochs=1):\n",
        "    gen_opt = optim.Adam(gen.parameters(), lr=2e-4)\n",
        "\n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "    gen_losses = []\n",
        "    steps = 0\n",
        "    for epoch in range(epochs):\n",
        "        for i, (real, _) in enumerate(loader):\n",
        "\n",
        "          noise = torch.randn(len(real), z_dim).to(device)\n",
        "          fake = gen(noise)\n",
        "\n",
        "          ### Train Generator ###\n",
        "          disc_fake_pred = disc(fake)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          gen_loss = get_gen_modified_BCE_loss(disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          gen_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          gen_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          gen_opt.step()\n",
        "\n",
        "          # update the loss trackers\n",
        "          gen_losses.append(gen_loss.item())\n",
        "          \n",
        "          # update the progress bar\n",
        "          steps += 1\n",
        "          progress_bar.update(progress(steps, gen_losses[-1], max_steps))\n",
        "    \n",
        "    plt.plot(gen_losses, label='Generator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "    return gen\n",
        "\n",
        "trained_gen = train_generator(disc_white_noise, gen, img_train_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2CwjEcFLsck3"
      },
      "source": [
        "Let's see what the generated images look like!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cg2hDerDsa7J"
      },
      "source": [
        "noise = torch.randn(24, z_dim).to(device)\n",
        "images = gen(noise).detach().cpu().view(24, C, H, W)\n",
        "imshow(make_grid(images, nrow=8))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QCnYfcn8Eq5x"
      },
      "source": [
        "## Step 3: Learn Discriminator with a Generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "aKx7Ekbhr9he"
      },
      "source": [
        "#@title Video: Training the Discriminator\n",
        "try: t4;\n",
        "except NameError: t4=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"y01DfHZ9_Vw\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GCIqvnLR1nkZ"
      },
      "source": [
        "We are going to freeze the generator and train the discriminator with it. \n",
        "\n",
        "Note that since the generator is needed when calculating the discriminator's loss, we need to call `.detach()` on the generator output (which removes the generator graph from the computational graph of discriminator loss) to ensure that only the discriminator gradients are calculated.\n",
        "\n",
        "In addition, using only `fake.detach()` in the discriminator loss computation allows us to use it later for the generator loss. Since, by default, after calling `.backward()` all the computational graphs used in the loss calculation will be removed, so if we do not use `fake.detach()`, we will not be able to backward pass through it again in generator loss backpropagation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KiW4sK3KtTdu"
      },
      "source": [
        " def train_discriminator_with_gen(disc, get, loader, epochs=1):\n",
        "    disc_opt = optim.Adam(disc.parameters(), lr=2e-4)\n",
        "\n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "    disc_losses = []\n",
        "    steps = 0\n",
        "    for epoch in range(epochs):\n",
        "        for i, (real, _) in enumerate(loader):\n",
        "          real = real.to(device)\n",
        "\n",
        "          # random noise for the generator\n",
        "          noise = torch.randn(len(real), z_dim).to(device)\n",
        "        \n",
        "          ### Train Discriminator ###\n",
        "          disc_real_pred = disc(real)\n",
        "          fake = gen(noise)\n",
        "          # Next don't forget to give a detached fake to the discriminator\n",
        "          # since we do not want to backdrop to generator yet \n",
        "          disc_fake_pred = disc(fake.detach())\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          disc_loss = get_disc_BCE_loss(disc_real_pred, disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          disc_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          disc_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          disc_opt.step()\n",
        "\n",
        "          # update the loss trackers\n",
        "          disc_losses.append(disc_loss.item())\n",
        "          \n",
        "          # update the progress bar\n",
        "          steps += 1\n",
        "          progress_bar.update(progress(steps, disc_losses[-1], max_steps))\n",
        "    \n",
        "    plt.plot(disc_losses, label='Discriminator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "    return disc\n",
        "\n",
        "disc_white_noise = train_discriminator_with_gen(disc_white_noise, trained_gen, img_train_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "voGgFnlYFqxI"
      },
      "source": [
        "## Step 4: Learn Discriminator and Generator together"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "e5KTlTzKu5-C"
      },
      "source": [
        "#@title Video: Full GAN\n",
        "try: t5;\n",
        "except NameError: t5=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"6GGdZ5eTNmQ\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XhZWrIz1_dMt"
      },
      "source": [
        "\n",
        "**Note** that we set $\\beta_1$ value of Adam optimizer, which deactivates the gradient's momentum. Why is momentum bad for GANs? Because the data distribution is non-stationary over training (since fake distribution is changing over time), the loss surfaces are kept changing after each update. Hence, a downhill might suddenly turn into an uphill, so sometimes there is no point in keep going in that direction; therefore, momentum, if used carelessly, often makes training more unstable!\n",
        "\n",
        "Now let's put everything together and see the magic!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8LgY1Q_U12jC"
      },
      "source": [
        " def train_full_GAN(gen, disc,\n",
        "                    loader, z_dim, \n",
        "                    epochs=5, disp_batch_size = 24):\n",
        "    gen.to(device).train()\n",
        "    disc.to(device).train()\n",
        "\n",
        "    disc_opt = optim.Adam(disc.parameters(), lr=2e-4, betas=(0.0, 0.999))\n",
        "    gen_opt = optim.Adam(gen.parameters(), lr=2e-4, betas=(0.0, 0.999))\n",
        "\n",
        "    # fix noise for testing genererator at each epoch\n",
        "    fix_noise = torch.randn(disp_batch_size, z_dim).to(device)\n",
        "\n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "    gen_losses = []\n",
        "    disc_losses = []\n",
        "    steps = 0\n",
        "    for epoch in range(epochs):\n",
        "        for i, (real, _) in enumerate(loader):\n",
        "          real = real.to(device)\n",
        "          batch_size = len(real)\n",
        "          \n",
        "          # random standard normal noise for generator\n",
        "          noise = torch.randn(batch_size, z_dim).to(device)\n",
        "\n",
        "          ### Train Discriminator ###\n",
        "          # Forward computations\n",
        "          fake = gen(noise)\n",
        "          # Next don't forget to give a detached fake to the discriminator\n",
        "          # since we do not want to backdrop to generator yet \n",
        "          disc_fake_pred = disc(fake.detach())\n",
        "          disc_real_pred = disc(real)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          disc_loss = get_disc_BCE_loss(disc_real_pred, disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          disc_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          disc_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          disc_opt.step()\n",
        "\n",
        "          ### Train Generator ###\n",
        "          # Forward computations\n",
        "          disc_fake_pred = disc(fake)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          gen_loss = get_gen_modified_BCE_loss(disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          gen_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          gen_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          gen_opt.step()\n",
        "\n",
        "          # update the loss trackers\n",
        "          gen_losses.append(gen_loss.item())\n",
        "          disc_losses.append(disc_loss.item())\n",
        "\n",
        "          progress_bar.update(progress(steps, (gen_losses[-1], disc_losses[-1]), max_steps))\n",
        "          steps += 1\n",
        "\n",
        "        images = gen(fix_noise).detach().cpu().view(disp_batch_size, C, H, W)\n",
        "        imshow(make_grid(images, nrow=8))\n",
        "    \n",
        "    gen.to('cpu').eval()\n",
        "    disc.to('cpu').eval()\n",
        "\n",
        "    plt.plot(gen_losses, label='Generator loss')\n",
        "    plt.plot(disc_losses, label='Discriminator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "\n",
        "gen = MLP(actv='PReLU()',\n",
        "          input_units = z_dim,\n",
        "          hidden_units_list = [128, 256],\n",
        "          output_units = C*H*W,\n",
        "          out_actv = 'Tanh()').to(device)\n",
        "\n",
        "disc = MLP(actv='PReLU()',\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [256, 128],\n",
        "          output_units = 1,\n",
        "          out_actv = 'Sigmoid()').to(device)\n",
        "\n",
        "train_full_GAN(gen, disc, img_train_loader, z_dim)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6DDgCxLXLYwP",
        "cellView": "form"
      },
      "source": [
        "#@markdown Let's assume the generator can always undoubtedly fool its discriminator. What would be the loss of this generator measured by the discriminator? from the training loss plot, is the final generator loss close to this value?\n",
        "w8_is_it_optimal_gen = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tkhl-_-yOvCG"
      },
      "source": [
        "**Answer**: Let's review the generator loss first,\n",
        "\n",
        "$$\n",
        "J_{G}=-\\frac{1}{m} \\sum_{i=1}^{fake} \\log \\left(D\\left(x_{i}\\right)\\right)\n",
        "$$\n",
        "\n",
        "Now imagine that all the fake examples are classified as real with 1.0 probability (since it is fooling undoubtedly). Therefore we have,\n",
        "$$\n",
        "J_{G}=-\\frac{1}{m} \\sum_{i=1}^{fake} \\log \\left(1.0\\right) = -\\log \\left(1.0\\right) = 0\n",
        "$$\n",
        "\n",
        "So we see that the minimum generator loss is 0. And the training loss is nowhere close to this value, indicating that the discriminator is far superior to the generator, and there is a danger of vanishing gradients!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oZWsGIbab_DI"
      },
      "source": [
        "---\n",
        "# Section 3: Convolutional GANs\n",
        "*Estimated time to here: 55min*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ZRISDk2ba-Ss"
      },
      "source": [
        "#@title Video: Convolutional GAN\n",
        "try: t6;\n",
        "except NameError: t6=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"IxveRWDe-yc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IPCi8mjnUPg3"
      },
      "source": [
        "Let's implement a general-purpose ConvNet for both discriminator and generator like the one with MLPs. Don't worry; The implementation is not an exercise, but you need to figure out how `ConvTranspose2d` is working to answer the next question."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kBzau0kZRlnq"
      },
      "source": [
        "class ConvNet(nn.Module):\n",
        "    def __init__(self, actv, input_channel, hidden_channels, out_actv=None, is_gen=False):\n",
        "        super(ConvNet, self).__init__()\n",
        "        self.is_gen = is_gen\n",
        "\n",
        "        # Initialize layers of MLP\n",
        "        self.layers = nn.ModuleList()\n",
        "\n",
        "        # Loop over layers and create each one\n",
        "        in_channels = input_channel\n",
        "        for i in range(len(hidden_channels)):\n",
        "          out_channels = hidden_channels[i] \n",
        "\n",
        "          if self.is_gen: # if it is a generator model\n",
        "            if i > 0:  # if first layer\n",
        "              self.layers += [nn.ConvTranspose2d(\n",
        "                                  in_channels=in_channels,\n",
        "                                  out_channels=out_channels,\n",
        "                                  kernel_size=4, #<--\n",
        "                                  stride=2, #<----\n",
        "                                  padding=1 #<---\n",
        "                              )]\n",
        "            else:\n",
        "                self.layers += [nn.ConvTranspose2d(\n",
        "                      in_channels=in_channels,\n",
        "                      out_channels=out_channels,\n",
        "                      kernel_size=2, #<--\n",
        "                      stride=1, #<----\n",
        "                      padding=0 #<---\n",
        "                  )]\n",
        "          else:\n",
        "            self.layers += [nn.Conv2d(\n",
        "                    in_channels=in_channels,\n",
        "                    out_channels=out_channels,\n",
        "                    kernel_size=4,\n",
        "                    stride=2,\n",
        "                    padding=1\n",
        "                )]\n",
        "\n",
        "          in_channels = out_channels\n",
        "          exec('self.activation = nn.%s'%actv)\n",
        "          self.layers += [\n",
        "            self.activation]  \n",
        "\n",
        "        # Create final layer\n",
        "        if self.is_gen:  \n",
        "          self.out = nn.ConvTranspose2d(\n",
        "                                in_channels=out_channels,\n",
        "                                out_channels=3,\n",
        "                                kernel_size=4,\n",
        "                                stride=2,\n",
        "                                padding=1,\n",
        "                            )\n",
        "        else:\n",
        "          self.out = nn.Conv2d(\n",
        "                      in_channels=out_channels,\n",
        "                      out_channels=1,\n",
        "                      kernel_size=2,\n",
        "                      stride=1,\n",
        "                      padding=0,\n",
        "                  )\n",
        "        \n",
        "        self.out_actv = out_actv\n",
        "        if out_actv:\n",
        "          exec('self.out_actv = nn.%s'%out_actv) \n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.is_gen: \n",
        "          # make inputs to 1x1 2D with channel noise\n",
        "          x = x.view(x.shape[0], -1, 1, 1)\n",
        "        else:\n",
        "          x = x.view(x.shape[0], -1, H, W)\n",
        "\n",
        "        # Get activations of each layer\n",
        "        for layer in self.layers:\n",
        "          x = layer(x)\n",
        "\n",
        "        # Get outputs\n",
        "        x = self.out(x)\n",
        "\n",
        "        if self.out_actv:\n",
        "          x = self.out_actv(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "U41neCflTfn7"
      },
      "source": [
        "#@markdown Why we use different kernel size and stride in the first ConvTranspose2d?\n",
        "w8_why_different_filter_in_first_layer_for_gen = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X48X9QomU4wf"
      },
      "source": [
        "**Answer**: In the first layer, the input has a 1x1 spatial dimension. And the goal of each layer is to double the spatial dimension. Hence we would like to achieve a 2x2 dimension in the second layer. We could achieve the same result by either option since the zero padding in 4x4 gets the center of 4x4 which is 2x2 (the stride does not matter here since the input is just 1x1). But a 2x2 is more efficient and a 4x4 filter is waseful. So that's it just for faster computations!\n",
        "\n",
        "\n",
        "<img align=\"center\" width=\"600\" height=\"200\" src=\"https://github.com/CIS-522/course-content/raw/main/tutorials/W08_AutoEncoders_GANs/static/deconv.jpg\">\n",
        "\n",
        "\n",
        "\n",
        "Note: I know this had nothing to do with GANs, but if you noticed this, you could really take pride in knowing your convolutions ;)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zdyxePZmlIrv"
      },
      "source": [
        "disc = ConvNet(actv='PReLU()',\n",
        "          input_channel = 3,\n",
        "          hidden_channels = [128, 256, 512, 1024], \n",
        "          out_actv = 'Sigmoid()', is_gen=False)\n",
        "\n",
        "gen = ConvNet(actv='PReLU()',\n",
        "          input_channel = z_dim,\n",
        "          hidden_channels = [1024, 512, 256, 128], \n",
        "          out_actv = 'Tanh()', is_gen=True)\n",
        "\n",
        "train_full_GAN(gen, disc, img_train_loader, z_dim)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "xL3MfrXFiOJ7"
      },
      "source": [
        "#@markdown The generator BCE loss is diverging, but images are still improving! Why do you think BCE loss fails to measure the actual performance of GANs?\n",
        "w8_why_bce_is_bad_metric = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VmxMXBcURfq9"
      },
      "source": [
        "**Answer**: Because there are many Nash equilibria, there is no optimal convergent point. So, as a result, the training losses often oscillate, and looking at only BCE loss trajectory is not enough to measure the model performance."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LWi9aCnHBGxQ"
      },
      "source": [
        "---\n",
        "# GAN Failures"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "HuBrUXyeA9kn"
      },
      "source": [
        "#@title Video: GAN Failure Modes\n",
        "try: t7;\n",
        "except NameError: t7=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"u-8mv1MHiec\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-FpTjA2gTkhc"
      },
      "source": [
        "## Exercise 3: Visualize Mode Collapse\n",
        "\n",
        "<img align=\"right\" width=\"300\" height=\"300\" src=\"https://raw.githubusercontent.com/CIS-522/course-content/main/tutorials/W08_AutoEncoders_GANs/static/ganlab.png\">\n",
        "\n",
        "Now that we have a good idea on what is Mode Collapse let's try to see it happen in real-time.\n",
        "\n",
        "\n",
        "For this exercise, you could follow these steps:\n",
        "1.   Familiarize yourself with the interface\n",
        "2.   In the Data Distribution tab to the right of GAN Lab logo, click on the last option with a yellow marker sign\n",
        "3.   Draw a distribution similar to the shown example (one pile being denser than others) and press APPLY\n",
        "4.   Press the play button to start training\n",
        "5.   Wait to see where the distribution generally lies\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "S86iu-mntya2"
      },
      "source": [
        "# @title GanLab\n",
        "from IPython.display import HTML\n",
        "HTML('<iframe width=\"1500\" height=\"660\" src=\"https://poloclub.github.io/ganlab/\" allowfullscreen></iframe>')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lzv7g54-c7gv"
      },
      "source": [
        "Reminder: The **Nash equilibrium** is a concept of game theory where the optimal outcome of a game is one where no player has an incentive to deviate from their chosen strategy after considering an opponent's choice."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "eiBYN4e5bJU9"
      },
      "source": [
        "#@markdown We saw that the generator distribution is going to converge at the big pile. Why do you think this Mode collapse is a desirable output for the game? Is it a Nash equilibrium?\n",
        "w8_is_it_Nash_equilibrium = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKnN1VXRjQv8"
      },
      "source": [
        "**Answer**: Since the discriminator does not have a criterion on batches of images (i.e., whether they cover all the modes), it is perfectly fooled if the generator can fit one of the modes. Therefore the generator does not desire to change its strategy once it can perfectly fit a mode. Also, the discriminator does not have a desire to reject samples from the overfit mode since on an individual basis they look identical to the real samples from that mode. Hence both agents don't have the desire to change, which makes it a Nash equilibrium."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rpk10dSa3pPS"
      },
      "source": [
        "---\n",
        "# Section 4: Conditional GANs\n",
        "*Estimated time to here: 75min*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "VJvavOtwDuxl"
      },
      "source": [
        "#@title Video: Conditional GANs\n",
        "try: t8;\n",
        "except NameError: t8=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"XqhbOsydKo0\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YyCoHucbt3sW"
      },
      "source": [
        "## Exercise 4: Augment Inputs with Labels\n",
        "\n",
        "<img align=\"right\" width=\"250\" height=\"200\" src=\"https://raw.githubusercontent.com/CIS-522/course-content/main/tutorials/W08_AutoEncoders_GANs/static/conditional.png\">\n",
        "\n",
        "To train a conditional GAN, we could extend the models' input to accept a one-hot encoding representation of the class labels. For the generator, we simply concatenate the noise vector with the one-hot class label. For the discriminator's input, we can extend the one-hot encoding of the class label to be the same size as the image, and then we could concatenate them as different channels. For example, in the figure, a picture of a cat with three RGB channels is extended with 1,0,0 one-hot encodings of the class label. Meaning that the first channel of the extended image would be all ones, and the second and third channels of the extended image would be all zeros.\n",
        "\n",
        "Now let's implement this operation. The code for noise augmentation is given bellow, so with an additional step, you can do the next one very similarly!\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NT2ADFZVp8Te"
      },
      "source": [
        "def augment_noise_with_labels(noise, labels):\n",
        "    '''\n",
        "    Function to return conditional input on the labels.\n",
        "\n",
        "    Parameters:\n",
        "        noise: noise tensor for generator, (batch_size, z_dim)\n",
        "        labels: the batch labels with 0 to C-1 values, (batch_size)\n",
        "\n",
        "    Returns:\n",
        "      extended input with one-hot labels, (batch_size, z_dim + num_classes)\n",
        "    '''          \n",
        "    onehot_z_target = torch.diag(torch.ones(num_classes))[labels].to(device)\n",
        "    augmented_noise = torch.cat([noise, onehot_z_target], axis=1)\n",
        "    return augmented_noise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZRDixKVsm6h4"
      },
      "source": [
        "Hint: You could use `F.interpolate` for extending a 1 x 1 input to H x W. \n",
        "\n",
        "`F.interpolate` Down/up samples the input to either the given `size` or the given `scale_factor`. The algorithm used for interpolation is determined by `mode`, which is the Nearest neighbor by default. The nearest neighbor interpolation is the desired mode here since we'd like to repeat the same values HxW times."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Uyq0oCHm8MG"
      },
      "source": [
        "def augment_image_with_labels(imgs, labels):\n",
        "    '''\n",
        "    Function to return conditional input on the labels.\n",
        "\n",
        "    Parameters:\n",
        "        img: img batch tensor for discriminator, (batch_size, C, H, W)\n",
        "        labels: the batch labels with 0 to C-1 values, (batch_size)\n",
        "\n",
        "    Returns:\n",
        "      extended input with one-hot labels, (batch_size, C + num_classes, H, W)\n",
        "    '''            \n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the augment image function\")\n",
        "    ####################################################################          \n",
        "    onehot_z_target = ...\n",
        "    onehot_channel_target = ...\n",
        "    augmented_img = ...\n",
        "    return augmented_img\n",
        "\n",
        "def augment_image_with_labels(img, labels):\n",
        "    '''\n",
        "    Function to return conditional input on the labels.\n",
        "\n",
        "    Parameters:\n",
        "        img: img batch tensor for discriminator, (batch_size, C, H, W)\n",
        "        labels: the batch labels with 0 to C-1 values, (batch_size)\n",
        "\n",
        "    Returns:\n",
        "      extended input with one-hot labels, (batch_size, C + num_classes, H, W)\n",
        "    '''                      \n",
        "    onehot_z_target = torch.diag(torch.ones(num_classes))[labels].to(device)\n",
        "    onehot_channel_target = F.interpolate(onehot_z_target.view(batch_size, -1, 1, 1), (H, W))\n",
        "    augmented_img = torch.cat([img, onehot_channel_target], axis=1)\n",
        "    return augmented_img"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7X18jTVb6EaG"
      },
      "source": [
        "Next, we just need to apply our function in three places where the noise input and real and fake images are generated, "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lA9wWDB83so9"
      },
      "source": [
        " def train_Conditional_GAN(gen, disc,\n",
        "                          loader, z_dim, \n",
        "                          epochs=5, disp_batch_size = 24):\n",
        "    gen.to(device).train()\n",
        "    disc.to(device).train()\n",
        "\n",
        "    disc_opt = optim.Adam(disc.parameters(), lr=2e-4, betas=(0.0, 0.999))\n",
        "    gen_opt = optim.Adam(gen.parameters(), lr=2e-4, betas=(0.0, 0.999))\n",
        "\n",
        "    max_steps = epochs*len(loader)\n",
        "    progress_bar = display(progress(0, 0, max_steps), display_id=True)\n",
        "    gen_losses = []\n",
        "    disc_losses = []\n",
        "    steps = 0\n",
        "\n",
        "    # create 3 rows of each class for visualization\n",
        "    fix_zs = torch.randn(disp_batch_size, z_dim).to(device).float()\n",
        "    class_zs = torch.zeros(disp_batch_size, num_classes).to(device).float()\n",
        "    class_zs[:8, 0] = 1 # one hot vector of first class\n",
        "    class_zs[8:16, 1] = 1 # of second class\n",
        "    class_zs[16:, 2] = 1 # of third class\n",
        "    fix_zs = torch.cat([fix_zs, class_zs], axis=1)\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        for i, (real, labels) in enumerate(loader):\n",
        "          labels = labels.to(device)\n",
        "          real = real.to(device)\n",
        "          batch_size = len(real)\n",
        "          \n",
        "          real = real.to(device)\n",
        "          batch_size = len(real)\n",
        "\n",
        "          # extend real images with labels\n",
        "          real = augment_image_with_labels(real, labels) ## new line ##\n",
        "\n",
        "          # random standard normal noise for generator\n",
        "          noise = torch.randn(batch_size, z_dim).to(device)\n",
        "\n",
        "          # extend generator input with labels\n",
        "          noise = augment_noise_with_labels(noise, labels) ## new line ##\n",
        "\n",
        "          ### Train Discriminator ###\n",
        "          # Forward computations\n",
        "          fake = gen(noise)\n",
        "\n",
        "          # extend fake images with labels\n",
        "          fake = augment_image_with_labels(fake, labels) ##  new line ##\n",
        "\n",
        "          # Next don't forget to give a detached fake to the discriminator\n",
        "          # since we do not want to backdrop to generator yet \n",
        "          disc_fake_pred = disc(fake.detach())\n",
        "          disc_real_pred = disc(real)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          disc_loss = get_disc_BCE_loss(disc_real_pred, disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          disc_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          disc_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          disc_opt.step()\n",
        "\n",
        "          ### Train Generator ###\n",
        "          # Forward computations\n",
        "          disc_fake_pred = disc(fake)\n",
        "\n",
        "          # Calculate discriminator loss\n",
        "          gen_loss = get_gen_modified_BCE_loss(disc_fake_pred)\n",
        "\n",
        "          # Zero out the gradients before backpropagation\n",
        "          gen_opt.zero_grad()\n",
        "\n",
        "          # Update gradients\n",
        "          gen_loss.backward()\n",
        "\n",
        "          # Update optimizer\n",
        "          gen_opt.step()\n",
        "\n",
        "          # update the loss trackers\n",
        "          gen_losses.append(gen_loss.item())\n",
        "          disc_losses.append(disc_loss.item())\n",
        "\n",
        "          progress_bar.update(progress(steps, (gen_losses[-1], disc_losses[-1]), max_steps))\n",
        "          steps += 1\n",
        "\n",
        "        images = gen(fix_zs).detach().cpu().view(disp_batch_size, C, H, W)\n",
        "        imshow(make_grid(images, nrow=8))\n",
        "    \n",
        "    gen.to('cpu').eval()\n",
        "    disc.to('cpu').eval()\n",
        "\n",
        "    plt.plot(gen_losses, label='Generator loss')\n",
        "    plt.plot(disc_losses, label='Discriminator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "### Uncomment below to test your function\n",
        "# Same settings as before with just different inputs\n",
        "disc = ConvNet(actv='PReLU()',\n",
        "          input_channel = 3 + num_classes, # note that we extend the model input size\n",
        "          hidden_channels = [128, 256, 512, 1024], \n",
        "          out_actv = 'Sigmoid()', is_gen=False)\n",
        "\n",
        "gen = ConvNet(actv='PReLU()',\n",
        "          input_channel = z_dim + num_classes, # note that we extend the model input size\n",
        "          hidden_channels = [1024, 512, 256, 128], \n",
        "          out_actv = 'Tanh()', is_gen=True)\n",
        "train_Conditional_GAN(gen, disc, img_train_loader, z_dim, epochs=5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "WhUrHg48GuwJ"
      },
      "source": [
        "#@markdown Comparing with regular GANs, Why do you think the conditional GAN is doing better?\n",
        "w8_why_conditional_is_better = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MKAzvCIs3gwx"
      },
      "source": [
        "**Answer**: The generated images look more distinguishable and diverse which goes to show that the mode collapse is ameliorated.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qt860ZuRxBi-"
      },
      "source": [
        "---\n",
        "# Section 5: Controllable GANs\n",
        "*Estimated time to here: 95min*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "quiLOha0ltep"
      },
      "source": [
        "#@title Video: Controllable GAN\n",
        "try: t9;\n",
        "except NameError: t9=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"BqQUMoHukIU\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EGSho5GBfKB-"
      },
      "source": [
        "## Exercise 5: BlonderGAN\n",
        "\n",
        "<img align=\"right\" width=\"500\" height=\"200\" src=\"https://raw.githubusercontent.com/CIS-522/course-content/main/tutorials/W08_AutoEncoders_GANs/static/blonder.png\">\n",
        "\n",
        "In this exercise, we will use a pre-trained face feature classifier to change the output of a pre-trained GAN in the direction we like. Specifically, we will change the hair colors to blond, hence the name BlonderGAN!\n",
        "\n",
        "\n",
        "The exercise part is simply figuring out how to update the noise vector to achieve the desired output, according to the given Figure. After backpropagation, we use the gradient of the noise vector, `noise.grad`, and using a `step_size`, we go one step in the direction of the gradient.\n",
        "\n",
        "Note that this is the first time we are using backprop to obtain the gradient of a variable rather than a parameter. The reason is that we aim to find the optimum input rather than the optimal model. Since this is a simpler task, we do not need to use advanced optimizers such as Adam since a vanilla SGD is good enough."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nH5ndAAIlzg3"
      },
      "source": [
        "def calculate_updated_noise(noise, step_size):\n",
        "    '''\n",
        "    Function to return noise vectors updated with stochastic gradient ascent.\n",
        "    Parameters:\n",
        "        noise: the current noise vectors. You have already called the backwards function on the target class\n",
        "          so you can access the gradient of the output class with respect to the noise by using noise.grad\n",
        "        step_size: the scalar amount by which you should weight the noise gradient\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the noise update function\")\n",
        "    ####################################################################    \n",
        "\n",
        "    new_noise = ...\n",
        "    return new_noise \n",
        "\n",
        "def calculate_updated_noise(noise, step_size):\n",
        "    '''\n",
        "    Function to return noise vectors updated with stochastic gradient ascent.\n",
        "    Parameters:\n",
        "        noise: the current noise vectors. You have already called the backwards function on the target class\n",
        "          so you can access the gradient of the output class with respect to the noise by using noise.grad\n",
        "        step_size: the scalar amount by which you should weight the noise gradient\n",
        "    '''\n",
        "    new_noise = noise + ( noise.grad * step_size)\n",
        "    return new_noise\n",
        "\n",
        "### Uncomment below to test your function\n",
        "noise = torch.ones(20, 20) * 3\n",
        "noise.requires_grad_()\n",
        "fake_classes = (noise ** 2).mean()\n",
        "fake_classes.backward()\n",
        "new_noise = calculate_updated_noise(noise, 0.1)\n",
        "test_value = torch.mean(new_noise)\n",
        "true_value = 3.001\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDdGjuwnp4xb"
      },
      "source": [
        "Now let's put it to the actual test and see the results on images while step by step changing the noise in the gradient direction. First we need to load the models,"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OYyEaJ_Aaqqs"
      },
      "source": [
        "# trained on high-quality celebrity faces \"celebA\" dataset\n",
        "# this model outputs 256 x 256 pixel images\n",
        "model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub',\n",
        "                      'PGAN', model_name='celebAHQ-256',\n",
        "                      pretrained=True, useGPU=device=='cuda')\n",
        "gen = model.netG.eval().to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4mwUgUuvqwN8"
      },
      "source": [
        "# try some random examples\n",
        "num_images = 4\n",
        "noise, _ = model.buildNoiseData(num_images)\n",
        "with torch.no_grad():\n",
        "    generated_images = torch.tanh(gen(noise))\n",
        "\n",
        "# let's plot these images using torchvision and matplotlib\n",
        "grid = make_grid(generated_images.clamp(min=-1, max=1), scale_each=True, normalize=True)\n",
        "plt.imshow(grid.permute(1, 2, 0).cpu().numpy())\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ojn_ZI_JhThs"
      },
      "source": [
        "Nice!. Finally, we got something that resembles life, perhaps on another planet!\n",
        "\n",
        "Now let's load the pre-trained feature classifier,"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BYrPmOiN7rSX"
      },
      "source": [
        "n_classes = 40\n",
        "classifier = Classifier(n_classes=n_classes).to(device)\n",
        "class_dict = torch.load(\"pretrained_classifier.pth\", map_location=torch.device(device))[\"classifier\"]\n",
        "classifier.load_state_dict(class_dict)\n",
        "classifier.eval()\n",
        "\n",
        "opt = torch.optim.Adam(classifier.parameters(), lr=0.01)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AJs56-1v83w2"
      },
      "source": [
        "# First generate a bunch of images with the generator\n",
        "n_images = 8\n",
        "fake_image_history = []\n",
        "grad_steps = 10 # Number of gradient steps to take\n",
        "skip = 2 # Number of gradient steps to skip in the visualization\n",
        "\n",
        "feature_names = [\"5oClockShadow\", \"ArchedEyebrows\", \"Attractive\", \"BagsUnderEyes\", \"Bald\", \"Bangs\",\n",
        "\"BigLips\", \"BigNose\", \"BlackHair\", \"BlondHair\", \"Blurry\", \"BrownHair\", \"BushyEyebrows\", \"Chubby\",\n",
        "\"DoubleChin\", \"Eyeglasses\", \"Goatee\", \"GrayHair\", \"HeavyMakeup\", \"HighCheekbones\", \"Male\", \n",
        "\"MouthSlightlyOpen\", \"Mustache\", \"NarrowEyes\", \"NoBeard\", \"OvalFace\", \"PaleSkin\", \"PointyNose\", \n",
        "\"RecedingHairline\", \"RosyCheeks\", \"Sideburn\", \"Smiling\", \"StraightHair\", \"WavyHair\", \"WearingEarrings\", \n",
        "\"WearingHat\", \"WearingLipstick\", \"WearingNecklace\", \"WearingNecktie\", \"Young\"]\n",
        "\n",
        "### Change me! ###\n",
        "# enter Bald for the question\n",
        "target_indices = feature_names.index(\"BlondHair\")\n",
        "\n",
        "# noise = get_noise(n_images, z_dim).to(device).requires_grad_()\n",
        "noise, _ = model.buildNoiseData(n_images)\n",
        "noise = noise.to(device).requires_grad_()\n",
        "for i in range(grad_steps):\n",
        "    opt.zero_grad()\n",
        "    # Resizing is required since the classsifier is trained with 64x64 images\n",
        "    fake = F.interpolate(torch.tanh(gen(noise)), size=(64, 64), mode='area')\n",
        "    fake_image_history += [fake]\n",
        "    fake_classes_score = classifier(fake)[:, target_indices].mean()\n",
        "    fake_classes_score.backward()\n",
        "    noise.data = calculate_updated_noise(noise, 1 / grad_steps)\n",
        "\n",
        "show_tensor_images(torch.cat(fake_image_history[::skip], dim=2), num_images=n_images, nrow=n_images)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ssuoeZvwtNkY"
      },
      "source": [
        "#@markdown Try with Bald feature; why it fails on women?\n",
        "w8_why_no_bald_women = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qvSWJvrZzxQC"
      },
      "source": [
        "**Answer**\n",
        "\n",
        "Since bald women's occurrence is very rare in the dataset, the generator did not learn a disentangled representation of baldness, and it thinks it is a male characteristic. Hence often trying to turn women into bald comes with other male features."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKXyKSd98uCG"
      },
      "source": [
        "---\n",
        "# Section 6: Cycle GANs\n",
        "*Estimated time to here: 115min*\n",
        "\n",
        "So far, we have been able only to generate an example of distribution. But what if we want to learn a particular mapping and do translation from one domain to another! We can not achieve this with a single GAN., so if one GAN is not enough, two should do the trick!\n",
        "\n",
        "Let's see how two GANs can learn the unsupervised unpaired translation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "PlS0yMQ9loHR"
      },
      "source": [
        "#@title Video: CycleGAN\n",
        "try: t10;\n",
        "except NameError: t10=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"d4pcWJ8-XtY\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "awJXahLVtSDH"
      },
      "source": [
        "Let's unscramble Week3 images by training a CycleGAN. And use the following AB notation for unpaired unsupervised translation and relate it to our task: \n",
        "\n",
        "A: source, scrambled images  \n",
        "B: target, original images\n",
        "\n",
        "And here are the generetators:\n",
        "\n",
        "$$\n",
        "G_{AB}(A) \\longrightarrow B \\\\\n",
        "\n",
        "G_{BA}(B) \\longrightarrow A\n",
        "$$\n",
        "\n",
        "And here are the discriminators:\n",
        "\n",
        "$$\n",
        "D_{A}(A) \\longrightarrow \\mathbf{P}(\\text{$A$ is real}) \\\\\n",
        "\n",
        "D_{B}(B) \\longrightarrow \\mathbf{P}(\\text{$B$ is real})\n",
        "$$\n",
        "\n",
        "So let's get started by first creeating the B dataloader which returnes scrambled images."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H-p1f_CCMnmX"
      },
      "source": [
        "def get_scrambled_data_loaders(batch_size, rand_perm):\n",
        "\n",
        "  # Note how we can create custom transforms in torchvision\n",
        "  class Scrambler():\n",
        "    def __call__(self, tensor):\n",
        "        tensor_shape = tensor.shape\n",
        "        tensor_flat = tensor.view(-1)\n",
        "        tensor_scrambled = tensor_flat[rand_perm]\n",
        "        return tensor_scrambled.view(tensor_shape)\n",
        "\n",
        "  # define the transform done only during training\n",
        "  augmentation_transforms = [transforms.RandomHorizontalFlip()]\n",
        "\n",
        "  # define the transform done in training and testing (after augmentation)\n",
        "  preprocessing_transforms = [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
        "  \n",
        "  # compose them together\n",
        "  train_transform = transforms.Compose(augmentation_transforms + preprocessing_transforms + [Scrambler()])\n",
        "  test_transform = transforms.Compose(preprocessing_transforms + [Scrambler()])\n",
        "\n",
        "  # using pathlib to be compatible with all OS's\n",
        "  data_path = pathlib.Path('.')/'afhq'\n",
        "\n",
        "  # define the dataset objects (they can load one by one)\n",
        "  img_train_dataset = ImageFolder(data_path/'train', transform=train_transform)\n",
        "  img_test_dataset = ImageFolder(data_path/'val', transform=test_transform)\n",
        "\n",
        "  # define the dataloader objects (they can load batch by batch)\n",
        "  img_train_loader = DataLoader(img_train_dataset, batch_size=batch_size, pin_memory=True,\n",
        "                                shuffle=True, num_workers=0, drop_last=True)\n",
        "  # num_workers can be set to higher if running on Colab Pro TPUs to speed up,\n",
        "  # with more than one worker, it will do multithreading to queue batches\n",
        "  img_test_loader = DataLoader(img_test_dataset, batch_size=batch_size,\n",
        "                          shuffle=False, num_workers=0)\n",
        "  \n",
        "  return img_train_loader, img_test_loader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9IT0VB3mNUQX"
      },
      "source": [
        "rand_perm = torch.randperm(C*H*W)\n",
        "reverse_rand_perm = torch.argsort(rand_perm)\n",
        "scrambled_img_train_loader, _ = get_scrambled_data_loaders(batch_size, rand_perm)\n",
        "img_train_loader, img_test_loader = get_data_loaders(batch_size)\n",
        "\n",
        "# get some random scrambled training images\n",
        "dataiter = iter(scrambled_img_train_loader)\n",
        "images, labels = dataiter.next()\n",
        "\n",
        "# checking whether we can reverse it images\n",
        "imshow(make_grid(images[:8], nrow=8))\n",
        "images_flat = images.view(batch_size, -1)\n",
        "unscrambled_images = images_flat[:, reverse_rand_perm].view(images.shape)\n",
        "imshow(make_grid(unscrambled_images[:8], nrow=8))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "99iSqv0z2GGM"
      },
      "source": [
        "## Exercise 6.1: Discriminator Loss\n",
        "This is the same as in the previous sections, but now using any (not just BCE) plug-and-play Pytorch loss function (loss functions in Pytorch are generally called criterion), so it will teach you a better and more efficient way of doing what you learned before, although less didactic.\n",
        "\n",
        "**Note**: The use of `torch.zeros_like` and `torch.ones_like` are highly encouraged for efficiency!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UXbG2yXy3I-e"
      },
      "source": [
        "def get_disc_loss(real_X, fake_X, disc_X, adv_criterion):\n",
        "    '''\n",
        "    Return the loss of the discriminator given inputs.\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        fake_X: the generated images of class X\n",
        "        disc_X: the discriminator for class X; takes images and returns real/fake class X\n",
        "            prediction matrices\n",
        "        adv_criterion: the adversarial loss function; takes the discriminator \n",
        "            predictions and the target labels and returns a adversarial \n",
        "            loss (which you aim to minimize)\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the disc loss function\")\n",
        "    ####################################################################    \n",
        "    # Don't forget to detach your generator!\n",
        "    disc_fake_X_pred = disc_X(fake_X.detach()) # Detach generator\n",
        "    disc_fake_X_loss = ...\n",
        "    disc_real_X_pred = disc_X(real_X)\n",
        "    disc_real_X_loss = ...\n",
        "    disc_loss = ...\n",
        "    return disc_loss\n",
        "\n",
        "def get_disc_loss(real_X, fake_X, disc_X, adv_criterion):\n",
        "    '''\n",
        "    Return the loss of the discriminator given inputs.\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        fake_X: the generated images of class X\n",
        "        disc_X: the discriminator for class X; takes images and returns real/fake class X\n",
        "            prediction matrices\n",
        "        adv_criterion: the adversarial loss function; takes the discriminator \n",
        "            predictions and the target labels and returns a adversarial \n",
        "            loss (which you aim to minimize)\n",
        "    '''\n",
        "    disc_fake_X_pred = disc_X(fake_X.detach()) # Detach generator\n",
        "    disc_fake_X_loss = adv_criterion(disc_fake_X_pred, torch.zeros_like(disc_fake_X_pred))\n",
        "    disc_real_X_pred = disc_X(real_X)\n",
        "    disc_real_X_loss = adv_criterion(disc_real_X_pred, torch.ones_like(disc_real_X_pred))\n",
        "    disc_loss = (disc_fake_X_loss + disc_real_X_loss) / 2\n",
        "    return disc_loss\n",
        "\n",
        "\n",
        "### Uncomment below to test your function\n",
        "test_disc_X = lambda x: x * 97\n",
        "test_real_X = torch.tensor(83.)\n",
        "test_fake_X = torch.tensor(89.)\n",
        "test_adv_criterion = lambda x, y: x * 79 + y * 73\n",
        "\n",
        "values = get_disc_loss(test_real_X, test_fake_X, test_disc_X, test_adv_criterion)\n",
        "test_value = torch.mean(values)\n",
        "true_value = 659054.500\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HGhLg2ZK7ES9"
      },
      "source": [
        "## Exercise 6.2: Generator Adversarial Loss\n",
        "The first component of the generator's loss you're going to implement is its adversarial loss—this once again is pretty similar to the GAN loss that you've implemented in the past. The important thing to note is that the criterion now is based on least-squares loss rather than binary cross-entropy loss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4dqv_jHx7zJ7"
      },
      "source": [
        "def get_gen_adversarial_loss(real_X, disc_Y, gen_XY, adv_criterion):\n",
        "    '''\n",
        "    Return the adversarial loss of the generator given inputs\n",
        "    (and the generated images for testing purposes).\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        disc_Y: the discriminator for class Y; takes images and returns real/fake class Y\n",
        "            prediction matrices\n",
        "        gen_XY: the generator for class X to Y; takes images and returns the images \n",
        "            transformed to class Y\n",
        "        adv_criterion: the adversarial loss function; takes the discriminator \n",
        "                  predictions and the target labels and returns a adversarial \n",
        "                  loss (which you aim to minimize)\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the gen adv loss function\")\n",
        "    ####################################################################  \n",
        "    fake_Y = ...\n",
        "    disc_fake_Y_pred = ...\n",
        "    adversarial_loss = ...\n",
        "    return adversarial_loss, fake_Y\n",
        "\n",
        "def get_gen_adversarial_loss(real_X, disc_Y, gen_XY, adv_criterion):\n",
        "    '''\n",
        "    Return the adversarial loss of the generator given inputs\n",
        "    (and the generated images for testing purposes).\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        disc_Y: the discriminator for class Y; takes images and returns real/fake class Y\n",
        "            prediction matrices\n",
        "        gen_XY: the generator for class X to Y; takes images and returns the images \n",
        "            transformed to class Y\n",
        "        adv_criterion: the adversarial loss function; takes the discriminator \n",
        "                  predictions and the target labels and returns a adversarial \n",
        "                  loss (which you aim to minimize)\n",
        "    '''\n",
        "    fake_Y = gen_XY(real_X)\n",
        "    disc_fake_Y_pred = disc_Y(fake_Y)\n",
        "    adversarial_loss = adv_criterion(disc_fake_Y_pred, torch.ones_like(disc_fake_Y_pred))\n",
        "    return adversarial_loss, fake_Y\n",
        "\n",
        "### Uncomment below to test your function\n",
        "test_disc_Y = lambda x: x * 97\n",
        "test_real_X = torch.tensor(83.)\n",
        "test_gen_XY = lambda x: x * 89\n",
        "test_adv_criterion = lambda x, y: x * 79 + y * 73\n",
        "values = get_gen_adversarial_loss(test_real_X, test_disc_Y, test_gen_XY, test_adv_criterion)[0]\n",
        "test_value = torch.mean(values)\n",
        "true_value = 56606652.000\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PigYg7b39C-s"
      },
      "source": [
        "## Exercise 6.3: Generator Cycle Consistency Loss\n",
        "\n",
        "Now, you can implement the final generator loss and the part that puts the \"cycle\" in CycleGAN: cycle consistency loss. This is used to ensure that when you put an image through one generator, that if it is then transformed back into the input class using the opposite generator, the image is the same as the original input image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hl14nupd9iFM"
      },
      "source": [
        "def get_cycle_consistency_loss(real_X, fake_Y, gen_YX, cycle_criterion):\n",
        "    '''\n",
        "    Return the cycle consistency loss of the generator given inputs\n",
        "    (and the generated images for testing purposes).\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        fake_Y: the generated images of class Y\n",
        "        gen_YX: the generator for class Y to X; takes images and returns the images \n",
        "            transformed to class X\n",
        "        cycle_criterion: the cycle consistency loss function; takes the real images from X and\n",
        "                        those images put through a X->Y generator and then Y->X generator\n",
        "                        and returns the cycle consistency loss (which you aim to minimize)\n",
        "    '''\n",
        "    ####################################################################\n",
        "    # Fill in all missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Define the cycle consistency function\")\n",
        "    ####################################################################   \n",
        "    cycle_X = ...\n",
        "    cycle_loss = ...\n",
        "    return cycle_loss, cycle_X\n",
        "\n",
        "\n",
        "def get_cycle_consistency_loss(real_X, fake_Y, gen_YX, cycle_criterion):\n",
        "    '''\n",
        "    Return the cycle consistency loss of the generator given inputs\n",
        "    (and the generated images for testing purposes).\n",
        "    Parameters:\n",
        "        real_X: the real images from pile X\n",
        "        fake_Y: the generated images of class Y\n",
        "        gen_YX: the generator for class Y to X; takes images and returns the images \n",
        "            transformed to class X\n",
        "        cycle_criterion: the cycle consistency loss function; takes the real images from X and\n",
        "                        those images put through a X->Y generator and then Y->X generator\n",
        "                        and returns the cycle consistency loss (which you aim to minimize)\n",
        "    '''\n",
        "    cycle_X = gen_YX(fake_Y)\n",
        "    cycle_loss = cycle_criterion(cycle_X, real_X)\n",
        "    return cycle_loss, cycle_X\n",
        "\n",
        "\n",
        "### Uncomment below to test your function\n",
        "test_real_X = torch.tensor(83.)\n",
        "test_fake_Y = torch.tensor(97.)\n",
        "test_gen_YX = lambda x: x * 89\n",
        "test_cycle_criterion = lambda x, y: (x + y) * 73\n",
        "values = get_cycle_consistency_loss(test_real_X, test_fake_Y, test_gen_YX, test_cycle_criterion)[0]\n",
        "test_value = torch.mean(values)\n",
        "true_value = 636268.000\n",
        "offset = abs(true_value - test_value.item())\n",
        "assert offset < 1e-3, \"Result: %.3f, Expected: %.3f\"%(test_value, true_value)\n",
        "print('Success!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vctqhIzq_yRk"
      },
      "source": [
        "For the sake of time, we skip the overall loss calculation exercise. Please read carefully to see how all the pieces you just wrote come together (especially if you decide to attempt the Kaggle competition for a second round)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rNME-xFckeOF"
      },
      "source": [
        "def get_gen_loss(real_A, real_B, gen_AB, gen_BA, disc_A, disc_B, adv_criterion,\n",
        "                 cycle_criterion, lambda_cycle=15):\n",
        "\n",
        "    # Hint 1: Make sure you include both directions - you can think of the generators as collaborating\n",
        "    # Hint 2: Don't forget to use the lambdas for the identity loss and cycle loss!\n",
        "    # Adversarial Loss -- get_gen_adversarial_loss(real_X, disc_Y, gen_XY, adv_criterion)\n",
        "    adv_loss_BA, fake_A = get_gen_adversarial_loss(real_B, disc_A, gen_BA, adv_criterion)\n",
        "    adv_loss_AB, fake_B = get_gen_adversarial_loss(real_A, disc_B, gen_AB, adv_criterion)\n",
        "    gen_adversarial_loss = adv_loss_BA + adv_loss_AB\n",
        "\n",
        "\n",
        "    # Cycle-consistency Loss -- get_cycle_consistency_loss(real_X, fake_Y, gen_YX, cycle_criterion)\n",
        "    cycle_loss_BA, cycle_A = get_cycle_consistency_loss(real_A, fake_B, gen_BA, cycle_criterion)\n",
        "    cycle_loss_AB, cycle_B = get_cycle_consistency_loss(real_B, fake_A, gen_AB, cycle_criterion)\n",
        "    gen_cycle_loss = cycle_loss_BA + cycle_loss_AB\n",
        "\n",
        "    # Total loss\n",
        "    gen_loss = lambda_cycle * gen_cycle_loss + gen_adversarial_loss\n",
        "    return gen_loss, fake_A, fake_B \n",
        " \n",
        "def train_unscrambler_GAN(gen_AB, gen_BA, gen_opt,\n",
        "                           disc_A, disc_B, disc_A_opt, disc_B_opt,\n",
        "                           loader_A, loader_B, epochs=3):\n",
        "    gen.to(device).train()\n",
        "    disc.to(device).train()\n",
        "\n",
        "    adv_criterion = nn.MSELoss()\n",
        "    recon_criterion = nn.L1Loss()\n",
        "\n",
        "    gen_losses = []\n",
        "    disc_losses = []\n",
        "    progress_bar = display(progress(0, 0, epochs), display_id=True)\n",
        "    for epoch in range(epochs):\n",
        "        for i, ((real_A, _), (real_B, _)) in enumerate(zip(loader_A, loader_B)):\n",
        "          real_A_shape = real_A.shape\n",
        "          batch_size = real_A_shape[0]\n",
        "          real_A = real_A.to(device).view(batch_size, -1)\n",
        "          real_B = real_B.to(device).view(batch_size, -1)\n",
        "\n",
        "          ### Update discriminator A ###\n",
        "          disc_A_opt.zero_grad() # Zero out the gradient before backpropagation\n",
        "          with torch.no_grad(): # since we don't need to backprop through gen_BA\n",
        "              fake_A = gen_BA(real_B)\n",
        "          disc_A_loss = get_disc_loss(real_A, fake_A, disc_A, adv_criterion)\n",
        "          disc_A_loss.backward() # Update gradients\n",
        "          disc_A_opt.step() # Update optimizer\n",
        "\n",
        "          ### Update discriminator B ###\n",
        "          disc_B_opt.zero_grad() # Zero out the gradient before backpropagation\n",
        "          with torch.no_grad(): # since we don't need to backprop through gen_AB\n",
        "              fake_B = gen_AB(real_A)\n",
        "          disc_B_loss = get_disc_loss(real_B, fake_B, disc_B, adv_criterion)\n",
        "          disc_B_loss.backward() # Update gradients\n",
        "          disc_B_opt.step() # Update optimizer\n",
        "\n",
        "          ### Update generator ###\n",
        "          gen_opt.zero_grad()\n",
        "          gen_loss, _, _ = get_gen_loss(\n",
        "              real_A, real_B, gen_AB, gen_BA, disc_A, disc_B, adv_criterion, recon_criterion\n",
        "          )\n",
        "          gen_loss.backward() # Update gradients\n",
        "          gen_opt.step() # Update optimizer\n",
        "\n",
        "          # Keep track of the average discriminator loss\n",
        "          mean_discriminator_loss = disc_A_loss.item() / batch_size\n",
        "          disc_losses += [mean_discriminator_loss]\n",
        "\n",
        "          # Keep track of the average generator loss\n",
        "          mean_generator_loss = gen_loss.item() / batch_size\n",
        "          gen_losses += [mean_generator_loss]\n",
        "\n",
        "          progress_bar.update(progress(epoch+1, (mean_generator_loss,\n",
        "                                                 mean_discriminator_loss), epochs))\n",
        "          \n",
        "        imshow(make_grid(fake_B.view(real_A_shape)[:8].detach().cpu(), nrow=8))\n",
        "\n",
        "        real_A_flat = real_A.view(batch_size, -1).detach().cpu()\n",
        "        real_A_unscambled = real_A_flat[:, reverse_rand_perm].view(real_A_shape)\n",
        "        imshow(make_grid(real_A_unscambled[:8], nrow=8))\n",
        "    \n",
        "    gen.to('cpu').eval()\n",
        "    disc.to('cpu').eval()\n",
        "\n",
        "    plt.plot(gen_losses, label='Generator loss')\n",
        "    plt.plot(disc_losses, label='Discriminator loss')\n",
        "    plt.xlabel('Batches')\n",
        "    plt.ylabel('Training loss')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "disc_A = MLP(actv='PReLU()',\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [2024, 1024, 512, 256],\n",
        "          output_units = 1,\n",
        "          out_actv = 'Sigmoid()').to(device)\n",
        "\n",
        "disc_B = MLP(actv='PReLU()',\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [2024, 1024, 512, 256],\n",
        "          output_units = 1,\n",
        "          out_actv = 'Sigmoid()').to(device)\n",
        "\n",
        "gen_AB = MLP(actv='PReLU()',\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [],\n",
        "          output_units = C*H*W,\n",
        "          out_actv = 'Tanh()').to(device)\n",
        "\n",
        "gen_BA = MLP(actv='PReLU()',\n",
        "          input_units = C*H*W,\n",
        "          hidden_units_list = [],\n",
        "          output_units = C*H*W,\n",
        "          out_actv = 'Tanh()').to(device)\n",
        "\n",
        "gen_parameters = list(gen_AB.parameters())+list(gen_BA.parameters())\n",
        "gen_opt = optim.Adam(gen_parameters, lr=5e-5, betas=(0.0, 0.999))\n",
        "disc_A_opt = optim.Adam(disc_A.parameters(), lr=5e-5, betas=(0.0, 0.999))\n",
        "disc_B_opt = optim.Adam(disc_B.parameters(), lr=5e-5, betas=(0.0, 0.999))\n",
        "\n",
        "### Uncomment below to train\n",
        "train_unscrambler_GAN(gen_AB, gen_BA, gen_opt,\n",
        "                      disc_A, disc_B, disc_A_opt, disc_B_opt,\n",
        "                      scrambled_img_train_loader, img_train_loader,\n",
        "                      epochs=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "6LkU15B6vy1h"
      },
      "source": [
        "#@markdown How can we improve the results?\n",
        "w8_next_step_in_unscrambling = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zDBJYw1-yp8c"
      },
      "source": [
        "**Answer:**\n",
        "\n",
        "Here are some ideas worth trying:\n",
        "\n",
        "* Make it conditional\n",
        "* Add a Sparse regularization to the generator\n",
        "* Better initialization of the generator function\n",
        "* Using an auto-encoder type for the generator\n",
        "* Different loss functions for reconstruction and adversarial losses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LT3RLL-TmLHg"
      },
      "source": [
        "---\n",
        "# Section 7: Wrap Up\n",
        "*Estimated time to here: 145min*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "SIecC03-mYu3"
      },
      "source": [
        "#@title Video: Wrap Up\n",
        "try: t11;\n",
        "except NameError: t11=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"AhAwuTZW13E\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3fWo9SxYEhrh"
      },
      "source": [
        "## Submit responses"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "oJ2tznn6EnPP"
      },
      "source": [
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "import time\n",
        "import numpy as np\n",
        "import urllib.parse\n",
        "from IPython.display import IFrame\n",
        "\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefill_fields = {}\n",
        "  for key in fields:\n",
        "      new_key = 'prefill_' + key\n",
        "      prefill_fields[new_key] = fields[key]\n",
        "  prefills = urllib.parse.urlencode(prefill_fields)\n",
        "  src = src + prefills\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "try: t6;\n",
        "except NameError: t6 = time.time()\n",
        "try: t7;\n",
        "except NameError: t7 = time.time()\n",
        "try: t8;\n",
        "except NameError: t8 = time.time()\n",
        "try: t9;\n",
        "except NameError: t9 = time.time()\n",
        "try: t10;\n",
        "except NameError: t10 = time.time()\n",
        "try: t11;\n",
        "except NameError: t11 = time.time()\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: w8_why_different_distribution;\n",
        "except NameError: w8_why_different_distribution = \"\"\n",
        "try: w8_does_it_work_with_another_white_noise;\n",
        "except NameError: w8_does_it_work_with_another_white_noise = \"\"\n",
        "try: w8_is_it_optimal_gen;\n",
        "except NameError: w8_is_it_optimal_gen = \"\"\n",
        "try: w8_why_different_filter_in_first_layer_for_gen;\n",
        "except NameError: w8_why_different_filter_in_first_layer_for_gen = \"\"\n",
        "try: w8_why_bce_is_bad_metric;\n",
        "except NameError: w8_why_bce_is_bad_metric = \"\"\n",
        "try: w8_is_it_Nash_equilibrium;\n",
        "except NameError: w8_is_it_Nash_equilibrium = \"\"\n",
        "try: w8_why_conditional_is_better;\n",
        "except NameError: w8_why_conditional_is_better = \"\"\n",
        "try: w8_why_no_bald_women;\n",
        "except NameError: w8_why_no_bald_women = \"\"\n",
        "try: w8_next_step_in_unscrambling;\n",
        "except NameError: w8_next_step_in_unscrambling = \"\"\n",
        "\n",
        "\n",
        "times = [(t-t0) for t in [t1,t2,t3,t4,t5,t6,t7,t8,t9,t10,t11]]\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"w8_why_different_distribution\": w8_why_different_distribution,\n",
        "          \"w8_does_it_work_with_another_white_noise\": w8_does_it_work_with_another_white_noise,\n",
        "          \"w8_is_it_optimal_gen\": w8_is_it_optimal_gen,\n",
        "          \"w8_why_different_filter_in_first_layer_for_gen\": w8_why_different_filter_in_first_layer_for_gen,\n",
        "          \"w8_why_bce_is_bad_metric\": w8_why_bce_is_bad_metric,\n",
        "          \"w8_is_it_Nash_equilibrium\": w8_is_it_Nash_equilibrium,\n",
        "          \"w8_why_conditional_is_better\": w8_why_conditional_is_better,\n",
        "          \"w8_why_no_bald_women\": w8_why_no_bald_women,\n",
        "          \"w8_next_step_in_unscrambling\": w8_next_step_in_unscrambling,\n",
        "          \"cumulative_times\": times}\n",
        "\n",
        "src = \"https://airtable.com/embed/shr6nKM1glOObVaoH?\"\n",
        "\n",
        "#now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9YCjQi49EqPQ"
      },
      "source": [
        "## Feedback\n",
        "How could this session have been better? How happy are you in your group? How do you feel right now?\n",
        "\n",
        "Feel free to use the embeded form below or use this link:\n",
        "<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://airtable.com/shrNSJ5ECXhNhsYss\">https://airtable.com/shrNSJ5ECXhNhsYss</a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vdc-J2tJEtdu"
      },
      "source": [
        "display(IFrame(src=\"https://airtable.com/embed/shrNSJ5ECXhNhsYss?backgroundColor=red\", width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}